import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from jqdatasdk import *
from typing import List, Tuple, Dict
import time


# 股票池过滤函数
def get_hs_a_stock_pool(date: str = None) -> List[str]:
    """
    获取沪深A股股票池（排除ST、退市、非普通股）
    :param date: 日期（格式'YYYY-MM-DD'），默认取当前日期前1个交易日
    :return: 符合条件的股票代码列表
    """
    # 1. 获取所有证券，筛选沪深A股（code以000、001、002、003、600、601、603、688开头） all_sec = get_all_securities(types=['stock'], date=date)
    # 筛选A股代码（排除B股、港股通等）
    hs_a_code = [code for code in all_sec.index if
                 code.startswith(('000', '001', '002', '003', '600', '601', '603', '688'))]
    print(hs_a_code)

    # 2. 排除ST股票和退市股票
    valid_codes = []
    for code in hs_a_code:
        # 检查是否为ST股票（聚宽is_st_stock接口，1次调用/股票）
        is_st = is_st_stock(code, date=date)
        # 检查是否退市（退市股票名称含"退"字）
        stock_name = all_sec.loc[code, 'display_name']
        if not is_st and "退" not in stock_name:
            valid_codes.append(code)

    print(f"沪深A股有效股票池：共{len(valid_codes)}只股票")
    return valid_codes


# 趋势与放量判断函数
def judge_trend_and_volume(code: str, N: int, M: float, end_date: str = None) -> Tuple[bool, bool, Dict]:
    """
    判断单只股票是否满足“近N天线性下跌”和“近1-2天放量”
    :param code: 股票代码（如'600519.XSHG'）
    :param N: 趋势判断周期（天）
    :param M: 放量倍数阈值（≥M倍平均成交量判定为放量）
    :param end_date: 数据结束日期（格式'YYYY-MM-DD'），默认取当前日期前1个交易日
    :return: (是否线性下跌, 是否放量, 详细数据字典)
    """
    # 1. 获取近N+1天K线数据（含收盘价、成交量，1次调用/股票）
    # 取N+1天是为了确保能计算近N天趋势（避免因非交易日导致数据不足）
    kline_df = get_price(
        security=code,
        end_date=end_date,
        count=N + 1,
        frequency='daily',
        felds=['close', 'volume'],
        skip_paused=True,  # 跳过停牌日
        fq='pre'  # 前复权，确保价格连续性
    )

    # 数据校验：若有效交易日不足N天，直接返回不满足
    if len(kline_df) < N:
        return False, False, {
            'code': code,
            'reason': f'近{N + 1}天有效交易日不足{N}天',
            'close_data': None,
            'volume_data': None,
            'trend_slope': None,
            'volume_multiple': None
        }

    # 截取近N天数据（排除最早1天，确保数据量为N天）
    kline_n_days = kline_df.iloc[-N:]
    close_series = kline_n_days['close'].values  # 近N天收盘价
    volume_series = kline_n_days['volume'].values  # 近N天成交量
    dates = kline_n_days.index.strftime('%Y-%m-%d').tolist()  # 近N天日期
    # 2. 判断近N天线性下跌趋势（线性回归斜率<0）
    X = np.arange(N).reshape(-1, 1)  # 自变量：时间序列（0,1,...,N-1）
    y = close_series  # 因变量：收盘价
    lr = LinearRegression()
    lr.ft(X, y)
    trend_slope = lr.coef_[0]  # 回归斜率（负：下跌，正：上涨）
    is_down_trend = trend_slope < 0

    # 3. 判断近1-2天是否放量
    is_vol_up = False
    volume_multiple = 0.0
    # 近1天成交量（最后1天）和近2天成交量（最后2天） recent_1_vol = volume_series[-1]
    recent_2_vol = volume_series[-2] if N >= 2 else 0

    # 计算“排除近1-2天之外的平均成交量”（需确保剩余天数≥1）
    if N <= 2:
        # 若N≤2，排除近1-2天后无剩余数据，用近N天平均成交量替代
        avg_vol = np.mean(volume_series[:-1]) if N > 1 else recent_2_vol
    else:
        avg_vol = np.mean(volume_series[:-2])  # 排除近2天，取前N-2天平均

    # 避免平均成交量为0（防止除零错误）
    avg_vol = max(avg_vol, 1)

    # 检查近1天或近2天是否放量（≥M倍平均成交量）
    multiple_1 = recent_1_vol / avg_vol  # 近1天成交量倍数
    multiple_2 = recent_2_vol / avg_vol if N >= 2 else 0  # 近2天成交量倍数

    if multiple_1 >= M or multiple_2 >= M:
        is_vol_up = True
        volume_multiple = max(multiple_1, multiple_2)  # 取最大倍数

    # 整理详细数据
    detail_data = {
        'code': code,
        'name': get_all_securities(types=['stock'], date=end_date).loc[code, 'display_name'], 'dates': dates,
        'close_series': close_series.round(2),
        'volume_series': volume_series,
        'trend_slope': trend_slope.round(4),
        'avg_vol_exclude_recent': avg_vol.round(0),
        'max_volume_multiple': volume_multiple.round(2),
        'is_down_trend': is_down_trend,
        'is_vol_up': is_vol_up
    }

    return is_down_trend, is_vol_up, detail_data


# 批量筛选与结果输出函数
def batch_screen_stocks(N: int, M: float, end_date: str = None, batch_size: int = 50) -> List[Dict]:
    """
    批量筛选沪深A股中满足条件的股票
    :param N: 趋势判断周期（天）
    :param M: 放量倍数阈值
    :param end_date: 数据结束日期
    :param batch_size: 批量处理大小（避免单次处理过多股票导致内存占用过高）
    :return: 满足条件的股票详细数据列表
    """
    # 1. 初始化统计变量（关键：统计数据调用次数）
    total_calls = 0  # 总数据调用次数
    valid_stock_count = 0  # 有效股票总数
    target_stock_count = 0  # 满足条件的股票总数
    result_list = []

    # 2. 获取股票池
    stock_pool = get_hs_a_stock_pool(date=end_date)
    valid_stock_count = len(stock_pool)
    # 累计调用次数：get_all_securities（1次） + is_st_stock（每只股票1次）
    total_calls += 1 + valid_stock_count

    # 3. 批量处理股票
    total_batches = (valid_stock_count + batch_size - 1) // batch_size  # 总批次数
    for batch_idx in range(total_batches):
        # 截取当前批次股票
        start_idx = batch_idx * batch_size
        end_idx = min((batch_idx + 1) * batch_size, valid_stock_count)
        current_batch = stock_pool[start_idx:end_idx]
        print(f"\n处理第{batch_idx + 1}/{total_batches}批，共{len(current_batch)}只股票...")

    # 遍历当前批次股票
    for code in current_batch:
        try:
            # 判断趋势与放量（1次get_price调用/股票）
            is_down, is_vol, detail = judge_trend_and_volume(code, N, M, end_date)
            total_calls += 1  # 累计get_price调用次数

            # 若满足条件，加入结果列表
            if is_down and is_vol:
                target_stock_count += 1
                result_list.append(detail)
                # 实时输出满足条件的股票信息（每行1次反馈，符合用户"每反馈一行算1次调用"的定义）
                print(
                    f"满足条件 #{target_stock_count}：{detail['code']}（{detail['name']}）- "f"下跌斜率：{detail['trend_slope']}，放量倍数：{detail['max_volume_multiple']}倍")

            # 防止调用频率过高，每处理10只股票休眠0.5秒（聚宽接口建议）
            if (current_batch.index(code) + 1) % 10 == 0:
                time.sleep(0.5)

        except Exception as e:
            # 捕获异常（如股票数据异常），不中断批量处理
            print(f"处理{code}时出错：{str(e)}")
            continue

    # 4. 输出最终统计信息
    print(f"\n" + "=" * 80)
    print(f"筛选完成！")
    print(f"1. 股票池规模：{valid_stock_count}只沪深A股（排除ST、退市）")
    print(f"2. 满足条件股票：{target_stock_count}只（近{N}天线性下跌 + 近1-2天放量≥{M}倍）")
    print(f"3. 总数据调用次数：{total_calls}次（含get_all_securities、is_st_stock、get_price）")
    print(f"4. 反馈行数（满足条件股票）：{target_stock_count}行（符合用户'每反馈一行算1次调用'定义）")
    print(f"=" * 80)

    return result_list


# 主函数（执行入口）
if __name__ == "__main__1":
    # 配置筛选参数（根据需求修改）
    N = 10  # 近10天线性下跌趋势
    M = 2.0  # 放量阈值：≥2倍平均成交量
    END_DATE = '2024-09-20'  # 数据结束日期（格式'YYYY-MM-DD'），默认取前1交易日可设为-None

    BATCH_SIZE = 50  # 批量处理大小（建议50-100，平衡效率与内存）
    # 执行批量筛选
    target_stocks = batch_screen_stocks(
        N=N,
        M=M,
        end_date=END_DATE,
        batch_size=BATCH_SIZE
    )

    # （可选）将结果保存为CSV文件（聚宽环境支持本地保存）
    if target_stocks:
        result_df = pd.DataFrame(target_stocks)
        # 选择关键列保存
        save_columns = [
            'code', 'name', 'dates', 'close_series', 'volume_series',
            'trend_slope', 'avg_vol_exclude_recent', 'max_volume_multiple'
        ]
        result_df[save_columns].to_csv(
            f"hs_a_down_vol_stocks_N{N}_M{M}_{END_DATE}.csv",
            index=False,
            encoding='utf-8-sig'
        )
        print(f"\n结果已保存至：hs_a_down_vol_stocks_N{N}_M{M}_{END_DATE}.csv")

if __name__ == "__main__":
    print('股票池过滤')
    get_hs_a_stock_pool()

